import numpy as np
import cv2
import base
import numpy
import scipy.ndimage
from imageio import imread
from numpy.ma.core import exp
from scipy.constants.constants import pi


myfloat = np.float64


def generate_ws(i, j, M, N):
    res = np.cos((i + 0.5 - N / 2) * np.pi / N)
    return res


def estws(map_ssim):
    N, M = map_ssim.shape
    ws_map = np.zeros_like(map_ssim)

    for i in range(N):
        for j in range(M):
            ws_map[i][j] = generate_ws(i, j, M, N)

    return ws_map
    # cv2.imwrite("ws_map.png",ws_map*255)
    # import pdb; pdb.set_trace()


def compute_ssim(img_mat_1, img_mat_2):
    # Variables for Gaussian kernel definition
    gaussian_kernel_sigma = 1.5
    gaussian_kernel_width = 11
    gaussian_kernel = numpy.zeros((gaussian_kernel_width, gaussian_kernel_width))

    # Fill Gaussian kernel
    for i in range(gaussian_kernel_width):
        for j in range(gaussian_kernel_width):
            gaussian_kernel[i, j] = (1 / (2 * pi * (gaussian_kernel_sigma**2))) * exp(
                -(((i - 5) ** 2) + ((j - 5) ** 2)) / (2 * (gaussian_kernel_sigma**2))
            )

    # Convert image matrices to double precision (like in the Matlab version)
    img_mat_1 = img_mat_1.astype(numpy.float64)
    img_mat_2 = img_mat_2.astype(numpy.float64)

    # Squares of input matrices
    img_mat_1_sq = img_mat_1**2
    img_mat_2_sq = img_mat_2**2
    img_mat_12 = img_mat_1 * img_mat_2

    # Means obtained by Gaussian filtering of inputs
    img_mat_mu_1 = scipy.ndimage.filters.convolve(img_mat_1, gaussian_kernel)
    img_mat_mu_2 = scipy.ndimage.filters.convolve(img_mat_2, gaussian_kernel)

    # Squares of means
    img_mat_mu_1_sq = img_mat_mu_1**2
    img_mat_mu_2_sq = img_mat_mu_2**2
    img_mat_mu_12 = img_mat_mu_1 * img_mat_mu_2

    # Variances obtained by Gaussian filtering of inputs' squares
    img_mat_sigma_1_sq = scipy.ndimage.filters.convolve(img_mat_1_sq, gaussian_kernel)
    img_mat_sigma_2_sq = scipy.ndimage.filters.convolve(img_mat_2_sq, gaussian_kernel)

    # Covariance
    img_mat_sigma_12 = scipy.ndimage.filters.convolve(img_mat_12, gaussian_kernel)

    # Centered squares of variances
    img_mat_sigma_1_sq = img_mat_sigma_1_sq - img_mat_mu_1_sq
    img_mat_sigma_2_sq = img_mat_sigma_2_sq - img_mat_mu_2_sq
    img_mat_sigma_12 = img_mat_sigma_12 - img_mat_mu_12

    # c1/c2 constants
    # First use: manual fitting
    c_1 = 6.5025
    c_2 = 58.5225

    # Second use: change k1,k2 & c1,c2 depend on L (width of color map)
    l = 255
    k_1 = 0.01
    c_1 = (k_1 * l) ** 2
    k_2 = 0.03
    c_2 = (k_2 * l) ** 2

    # Numerator of SSIM
    num_ssim = (2 * img_mat_mu_12 + c_1) * (2 * img_mat_sigma_12 + c_2)
    # Denominator of SSIM
    den_ssim = (img_mat_mu_1_sq + img_mat_mu_2_sq + c_1) * (
        img_mat_sigma_1_sq + img_mat_sigma_2_sq + c_2
    )
    # SSIM
    ssim_map = num_ssim / den_ssim
    index = numpy.average(ssim_map)

    # print index

    return ssim_map, index


def _ws_ssim(image1, image2):
    map_ssim, MSSIM = compute_ssim(image1, image2)
    ws = estws(map_ssim)
    wsssim = np.sum(map_ssim * ws) / ws.sum()
    # print(wsssim)
    # print("WS-SSIM ",wsssim)

    return wsssim


def ws_ssim(img1, img2, mode="RGB"):
    if not img1.shape == img2.shape:
        raise ValueError('输入图片不相同')

    assert mode in ["RGB", "Y"], "运行模式只能是 RGB 或 Y"

    h, w, c = np.shape(img1)

    if c == 3 and mode == "Y":
        img1, img2 = base.rgb2y(img1, img2)

    return _ws_ssim(img1, img2)
